# -*- coding: utf-8 -*-
"""
@file name      : 1_cuda_methods.py
@author         : QuZhang
@date           : 2020-1-2
@brief          : 数据迁移到cuda的方法
"""
import torch
import torch.nn as nn


if __name__ == "__main__":
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    # ------------- tensor to cuda
    # 将数据迁移到gpu
    flag = True
    if flag:
        x_cpu = torch.ones((3, 3))
        print("x_cpu:\ndevice: {} is_cuda: {} id: {}".format(x_cpu.device, x_cpu.is_cuda, id(x_cpu)))

        x_gpu = x_cpu.to(device)  # 数据需要重新赋值，返回在device上的对象
        print("x_gpu:\ndevice：{} is_cuda: {} id: {}".format(x_gpu.device, x_gpu.is_cuda, id(x_gpu)))

    flag = True
    if flag:
        net = nn.Sequential(nn.Linear(3, 3))
        print("\nnet id:{} is_cuda: {}".format(id(net), next(net.parameters()).is_cuda))

        net.to(device)  # 模型不需要重新赋值
        print("net id:{} is_cuda: {}".format(id(net), next(net.parameters()).is_cuda))

    # ---------- forward in cuda
    # 在GPU上进行计算,模型和数据都要在GPU上
    flag = True
    if flag:
        output = net(x_gpu)
        print("output is_cuda: {}".format(output.is_cuda))  # 在GPU上运算的结果还是保持到GPU上

        # output = net(x_cpu)  # 一个cpu，一个gpu就会出错
